from functools import reduce
import math
import tensorflow as tf
import tensorflow.keras as keras
import tensorflow.keras.backend as K
from tensorflow.keras import layers
from losses_utils import *


# @tf.function
def yolo_loss_on_size(labels,predns,grid,grid_anchors,
                      nc=80,na=3,label_smoothing=0.0001,positive_weight=0.05,
                      xyxy=False,GIoU=False,DIoU=False,CIoU=False,
                      filter=True,weight=False):
    """
        predns为张量, (B,nx*ny,3,5+nc)
        labels为张量, (B,nx*ny,3,5+nc) [xc,yx,w,h,r,c...]
        #! loop 采用for循环计算, 还是张量整体计算
        #! filter 将正样本分离出来计算, 还是不分离计算
        #! weight 按正样本数量计算, 还是按权重计算
    """
    if not isinstance(labels,tf.Tensor): labels = tf.convert_to_tensor(labels,dtype=predns.dtype)
    else: labels = tf.cast(labels,dtype=predns.dtype)

    #! 初始化变量
    lbox = 0.0
    lbxy = 0.0
    lbwh = 0.0
    lcls = 0.0
    lobj = 0.0

    #! 还原pbox 从 偏移量 到 具体坐标
    pxy =  tf.nn.sigmoid(predns[:,:,:,0:2]) * 2  -  0.5     # (B,ny*nx,na,2) + (ny*nx,na,2)
    pwh = (tf.nn.sigmoid(predns[:,:,:,2:4]) * 2) ** 2 * tf.cast(grid_anchors,dtype=predns.dtype) # (B,ny*nx,na,2) * (na,2)
    predns = tf.concat([pxy,pwh,predns[...,4:]],axis=-1)

    condition = tf.equal(labels[...,4],1.0) #! (B,nx*ny,3)
    indices   = tf.where(condition)
    if tf.reduce_any(condition): #! 如果存在匹配
        t = tf.gather_nd(labels,tf.where(condition)) #! (N,nc)
        p = tf.gather_nd(predns,tf.where(condition)) #! (N,nc)

        iou = bbox_iou(t[...,:4],p[...,:4],xyxy=xyxy,GIoU=GIoU,DIoU=DIoU,CIoU=CIoU)
        lbox = tf.reduce_mean(1-iou)

        lbxy = tf.reduce_mean(tf.reduce_sum((t[...,:2] - p[...,:2])**2,axis=-1)) #! MSE
        lbwh = tf.reduce_mean(tf.reduce_sum(
                tf.abs(tf.nn.tanh(tf.math.log(tf.maximum(p[...,2:4] / t[...,2:4],K.epsilon())))),
                axis=-1
            ))

        # lcls = FocalLoss(BCEWithLogitsLoss(reduction='none'),reduction='mean')(
        #     (1-label_smoothing) * t[...,5:] + label_smoothing / nc,
        #     p[...,5:],
        # )

        lcls = BCEWithLogitsLoss(reduction='mean')(
            (1-label_smoothing) * t[...,5:] + label_smoothing / nc,
            p[...,5:],
        )
       
        tobj = tf.where( #! (B,nx*ny,na)
            condition,
            tf.maximum(bbox_iou(labels[...,:4],predns[...,:4],xyxy=xyxy,GIoU=GIoU,DIoU=DIoU,CIoU=CIoU),0.0),
            tf.zeros_like(labels[...,4])
        )

        lobj = FocalLoss(BCEWithLogitsLoss(pos_weight=1/positive_weight,reduction='none'),alpha=positive_weight,reduction='mean')(
            tobj[...,None],
            predns[...,4:5],
        )   
        
        # lobj = MyLoss(pos_weight=positive_weight,reduction='mean')(
        #     tobj[...,None],
        #     predns[...,4:5],
        # )
        
        # lobj = tf.reduce_mean(keras.losses.binary_crossentropy(
        #     tobj[...,None],
        #     predns[...,4:5],
        #     from_logits=True
        # ))

    return lbox, lbxy, lbwh, lcls, lobj

@tf.function
def yolo_loss(labels,predns,grids,grid_anchors,
              nc=80,na=3,label_smoothing=0.0001,positive_weight=0.05,
              xyxy=False,GIoU=False,DIoU=False,CIoU=False,
              filter=True,weight=False):
    """
        predns为张量, (B,nx*ny,3,5+nc)
        labels为张量, (B,nx*ny,3,5+nc) [xc,yx,w,h,r,c...]
        #! loop为False的情况下, filter失效
        #! loop为True, filter为True的情况下, weight失效
    """
    assert len(labels) == len(predns)

    lbox = 0.0
    lbxy = 0.0
    lbwh = 0.0
    lcls = 0.0
    lobj = 0.0
    # obj_balance = [4.0 , 1.0, 0.4 ]
    # bxy_balance = [4.0 , 1.0, 0.25]
    obj_balance = [0.4 , 1.0, 4.0]
    bxy_balance = [0.25, 1.0, 4.0]

    for l,(label,predn) in enumerate(zip(labels,predns)):
        box_loss, bxy_loss, bwh_loss, cls_loss, obj_loss = yolo_loss_on_size(
            label,predn,grids[l],grid_anchors[l],
            nc=nc,na=na,label_smoothing=label_smoothing,positive_weight=positive_weight,
            xyxy=xyxy,GIoU=GIoU,DIoU=DIoU,CIoU=CIoU,
            filter=filter,weight=weight)
        lbox += box_loss
        lbxy += bxy_loss
        lbxy += bxy_loss * bxy_balance[l]
        lbwh += bwh_loss
        lcls += cls_loss
        lobj += obj_loss * obj_balance[l]

    #! 使用model.compile中的loss_fn功能会把损失进行reduce化, 默认reduce函数是reduce_mean
    #! 可以使用reduce_sum, reduce_none, reduce_sum_over_batch进行替换, 但需要使用类处理
    #! 若直接使用loss_fn传入model.compile, 则会直接reduce_auto(默认是reduce_mean)
    
    return lbox, lbxy, lbwh, lcls, lobj


class YoloLoss(object):
    def __init__(self,grids,grid_anchors,nc=80,na=3,label_smoothing=0.0001,positive_weight=0.05,
                 xyxy=False,GIoU=False,DIoU=False,CIoU=False,filter=True,weight=False):
        self.grids = grids
        self.grid_anchors = grid_anchors
        self.nc = nc
        self.na = na
        self.label_smoothing = label_smoothing
        self.positive_weight = positive_weight
        self.xyxy = xyxy
        self.GIoU = GIoU
        self.DIoU = DIoU
        self.CIoU = CIoU
        self.filter = filter
        self.weight = weight
        self.obj_balance = [0.4 , 1.0, 4.0]
        self.bxy_balance = [0.25, 1.0, 4.0]
    
    # @tf.function
    def __call__(self,labels,predns):
        assert len(labels) == len(predns)

        lbox = 0.0
        lbxy = 0.0
        lbwh = 0.0
        lcls = 0.0
        lobj = 0.0
        # obj_balance = [4.0 , 1.0, 0.4 ]
        # bxy_balance = [4.0 , 1.0, 0.25]
        # obj_balance = [0.4 , 1.0, 4.0]
        # bxy_balance = [0.25, 1.0, 4.0]

        for l,(_labels,_predns) in enumerate(zip(labels,predns)):
            _lbox, _lbxy, _lbwh, _lcls, _lobj = yolo_loss_on_size(
                _labels,_predns,self.grids[l],self.grid_anchors[l],
                nc=self.nc,na=self.na,label_smoothing=self.label_smoothing,positive_weight=self.positive_weight,
                xyxy=self.xyxy,GIoU=self.GIoU,DIoU=self.DIoU,CIoU=self.CIoU,
                filter=self.filter,weight=self.weight)
            
            tf.print("l: ",l)
            tf.print("_predns shape: ",_predns.shape)
            tf.print("balance[l]: ",self.obj_balance[l])
            tf.print("lbox: ",_lbox)
            tf.print("lcls: ",_lcls)
            tf.print("lobj: ",_lobj)
                   
            lbox += _lbox
            lbxy += _lbxy * self.bxy_balance[l]
            lbwh += _lbwh
            lcls += _lcls
            lobj += _lobj * self.obj_balance[l]
        
        return lbox, lbxy, lbwh, lcls, lobj